% This function updates the (wage, ability) distribution and the endogenous
% grid over wages
function Probm = update_dist_int_edu_fix_p(Prob, P_S, Pr, Ss, rs, flag)
    global bet0 bet1 tau a b w e pi_theta N_a N_theta N_w N_e w_max w_min alpha xi xi1 ub lb pi_e...
        w theta Int_pi_w sigmae elasts log_e_mu log_e_sigma phi sigma gamma rts w_mult wspread frac0...
        unemp skill_shock
    
    Probm = zeros(size(Prob));
    ln_e_hat = -1/2*log_e_sigma^2;
    
    [~, F_es, wprimeS, ~] = choiceProb(Ss, rs);
    if flag == 1
        F_es = F_es*skill_shock;
    end

    % Wage grid is endogenous so that the approximation stays accurate when
    % wage growth occurs
    myDist1  = sum(bsxfun(@times, Prob .* P_S , reshape(pi_theta, [1,1,1,2])),4);
    w_bar_p = sum(myDist1 .* wprimeS, 1:3);
    
    w2 = w;     
    for t = 1:2
        for j = 1:N_a
            for i = 1:N_w
                wage = w(i);
                
                % The wage shock is i.i.d. and lognormally distributed
                % Take midpoints and allocate mass to closest wage grid
                % point
                t1 = unemp + (b   + squeeze(F_es(i,j,:))')  .* f_w(wage);
                midpts = log(w2) - log(t1);
                pts = ([midpts(1,:);(midpts(2:end,:) + midpts(1:end-1,:))/2; Inf Inf Inf] - ln_e_hat)/log_e_sigma;
                mean_prob = diff(normcdf(pts));
                tmp = bsxfun(@times, mean_prob*squeeze(P_S(i,j,:,t)), Pr(j,:));
                pr1 = Prob(i, j) .*  tmp .* pi_theta(t);
                Probm(:, :) = Probm(:, :) + pr1;
          
            end
        end
       
    end
    Probm = Probm/sum(Probm(:));
    

end